import copy
import logging
import time
from multiprocessing import Value, Manager

import cv2

import cjc.sort_path
import cjc.yolo_path

from cjc import project_setting
from cjc.project_setting import media_ext
from cjc.tools.logger import Logger
from cjc.tools.save_video import Cv2StreamSaver

from cjc.tracker.annotator import Annotator
from cjc.tracker.annotator_plot import AnnotatorP
from cjc.tracker.cls_counter import PersonCounter
from cjc.tracker.detector_for_track import DetectorForTrack
from cjc.tracker.detector_scheduler import DetectorScheduler
from cjc.tracker.opt_tool import create_opt
from cjc.tracker.save_one_box import crop_one_box
from cjc.tracker.sorter_for_track import Sorter
from cjc.tracker.tracker_tool import put_tracker_data_2
from cjc.tracker.traj_map import Trajectories
from cjc.detector.detect_age import AgeDetector
from cjc.detector.detect_sex import SexDetector
from cjc.detector.instance import InstanceDetector, create_instance_opt


class Tracker:

    def __init__(self, _opt):
        self.when_first_person_enter = None  #
        self.log = Logger(__class__.__name__)
        self.can_track = None
        self.tracker_queue = None  # 暂定None，需要传入
        self.need_to_refresh_path = None  #
        self.has_data_to_be_saved = False  #
        self.opt = _opt
        tmp = copy.deepcopy(_opt)
        tmp.source = project_setting.camera_index
        tmp.classes = 0

        self.detector = DetectorForTrack(_opt=tmp,
                                         update_sort_ages_cb=self.update_sort_ages_cb,
                                         update_camera_cb=self.update_camera_cb,
                                         update_current_frame_cb=self.update_current_frame_cb,
                                         pass_detections_to_sort_cb=self.pass_detections_to_sort_cb,
                                         sort_swap_frames_cb=self.sort_swap_frames_cb,
                                         nothing_detected_cb=self.nothing_detected_cb)
        self.traj_map = Trajectories()  # 轨迹图
        self.person_counter = PersonCounter(None)
        self.scheduler = DetectorScheduler(10)

        detector_opt = copy.deepcopy(_opt)
        detector_opt.view_img = False
        # sex_detector = SexDetector(detector_opt)
        # age_detector = AgeDetector(detector_opt)
        instance_detector = InstanceDetector(create_instance_opt())
        # self.scheduler.add_detector(sex_detector)
        # self.scheduler.add_detector(age_detector)
        # self.scheduler.add_detector(instance_detector)

        self.vid_cap = None  # 录制的时候才产生这个
        fp = f'{_opt.project}/webcam' + media_ext
        self.video_saver = Cv2StreamSaver(self.vid_cap, fp)
        self.video_saver.refresh_path()
        self.last_video_save_path = self.video_saver.save_path
        self.new_record_delay = 3  # 3后刷新`

        # self.annotator=None
        self.colors = self.detector.colors
        self.sorter = Sorter(_opt.device, self.detector.webcam, _opt.config_strongsort, _opt.strong_sort_weights)

    def sort_swap_frames_cb(self, i, im0):
        """保存相关 最后交换frame"""
        # self.log.info("sort_swipe_frames_cb",type(im0))
        if self.can_track.value == 1:
            #  如果 can_track 不可用 则跳过
            if self.has_data_to_be_saved:  # 正常
                self.traj_map.set_background(im0)
                self.video_saver.refresh_path()  # 刷新路径，结束当前录制  一定要先刷新在做其他的。
                img = self.traj_map.result_to_tracker()
                self.need_to_refresh_path = False  # 刚刷新完不用再进来了
                self.has_data_to_be_saved = False  # 该保存的保存完了
                self.update_first_person_enter()  # 3秒之内不在回来
                self.last_video_save_path = self.video_saver.save_path
                put_tracker_data_2(self.tracker_queue, img,
                                   self.person_counter.get_counter_status_for_webapi(),
                                   self.last_video_save_path, True)
                self.log.info("加入隊列保存", self.last_video_save_path)
                self.person_counter.clear()  # 前面已经把结果传入了ws 现在清空缓存
                self.traj_map.clear()  # 清楚轨迹图
            else:
                # pass
                img = self.traj_map.combine_traj_and_img(im0)
                # self.log.info("get_counter_status", self.person_counter.get_counter_status())
                put_tracker_data_2(self.tracker_queue, img, self.person_counter.get_counter_status(),
                                   self.last_video_save_path)
        if len(self.person_counter.pos_dic) > 0:
            # self.log.info("im0.shape.len",len(im0.shape))
            # self.log.info("im0.shape.len", im0[0])
            img = self.traj_map.uint8_result()
            self.video_saver.write(img)
        self.sorter.prev_frames[i] = self.sorter.curr_frames[i]

    def nothing_detected_cb(self, id):
        """ 准备刷新文件名 True，标记有数据需要保存  True """
        # self.sorter.strongsort_list[i].increment_ages()
        if self.has_data_to_be_saved is False and self.need_to_refresh_path:
            if self.how_long_first_person_entered() > self.new_record_delay:
                self.log.warning("超过了3秒")
                self.has_data_to_be_saved = True
                self.need_to_refresh_path = True

    def add_detect_result(self, _id, sex_cls, age_cls):
        if _id is None:
            print("id is None", _id)
            return
        if sex_cls is not None:
            self.person_counter.sex_group.add_person(int(_id), int(sex_cls))
        if age_cls is not None:
            self.person_counter.age_group.add_person(int(_id), int(age_cls))

    def update_sort_ages_cb(self, i):
        self.sorter.strongsort_list[i].increment_ages()

    def update_camera_cb(self, i):  # 相机运动补偿回调
        if self.sorter.cfg.STRONGSORT.ECC:  # camera motion compensation
            self.sorter.strongsort_list[i].tracker.camera_update(self.sorter.prev_frames[i], self.sorter.curr_frames[i])

    def update_current_frame_cb(self, i, im0):
        self.sorter.curr_frames[i] = im0
        self.traj_map.create_canvas_from_bk(im0)  # 设置 traj_map的背景图

    def pass_detections_to_sort_cb(self, i, xywhs, confs, clss, im0):
        self.person_counter.set_mask(im0)
        # self.log.warning("can_track",self.can_track.value)
        if self.can_track.value == 0:
            return  # 如果 can_track 不可用 则跳过
        # pass detections to strongsort
        self.sorter.outputs[i] = self.sorter.strongsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
        # draw boxes for visualization

        annotator = AnnotatorP(im0, line_thickness=1)
        # self.log.info("self.sorter.outputs[i]",self.sorter.outputs[i]
        if len(self.sorter.outputs[i]) > 0:
            self.update_first_person_enter()
            self.need_to_refresh_path = True  # 可以刷新了，但是现在还没到时候
            for j, (output, conf) in enumerate(zip(self.sorter.outputs[i], confs)):
                bboxes = output[0:4]  # 这个是检测结果，包括位置信息
                id = output[4]  # 跟踪结果id
                cls = output[5]  # 分类
                self.traj_map.draw_point_with_index(bboxes, id)  # 根据id设置颜色
                self.person_counter.run(id, bboxes)
                c = int(cls)  # integer class
                id = int(id)  # integer id
                imc = im0
                crop = crop_one_box(bboxes, imc, BGR=True)
                # 检测性别和年龄
                sex_finished, age_finished = self.person_counter.check_sex_and_age_detection_finished(id)
                detector_name = self.scheduler.get_current_detector_name()
                if detector_name == 'sex' and sex_finished:
                    # 跳过检测
                    pass
                elif detector_name == 'age' and age_finished:
                    # 跳过检测
                    pass
                else:
                    self.scheduler.detect(crop, id)

                sex_cls, sex_txt = self.scheduler.get_result_cls_name('sex')
                age_cls, age_txt = self.scheduler.get_result_cls_name('age')
                if sex_txt is None and age_txt is None:
                    label = f'{id} person '  #
                else:
                    label = f'{id}  {sex_txt} {age_txt} '  #
                annotator.box_label(bboxes, label, self.colors[c])
                self.add_detect_result(id, sex_cls, age_cls)  # age_cls
        outimg = annotator.result()
        if self.opt.view_img:
            cv2.imshow(__class__.__name__, outimg)
            cv2.waitKey(12)
            # cv2.imshow('tracker', im0)

    def set_conn(self, track_queue, can_track):
        """设置共享内存 queue 和 int"""
        if track_queue is not None:
            self.tracker_queue = track_queue
        if can_track is not None:
            self.can_track = can_track

    def update_first_person_enter(self):
        """第一个人进入后进入后计时三秒钟"""
        self.when_first_person_enter = time.time()

    def how_long_first_person_entered(self):
        return time.time() - self.when_first_person_enter

    def run(self):
        self.detector.detect()
        pass


def processor_run_tracker(track_queue, can_track):
    _opt = create_opt(
        yolo_weights='./weights/yolov7-tiny.pt'
        , strong_sort_weights='./weights/osnet_x0_25_msmt17.pt'
        , config_strongsort='./strong_sort/configs/strong_sort.yaml'
        , project='./runs/track',
        source=project_setting.camera_index
    )
    _opt.show_vid = project_setting.show_track_window  # 不显示
    tracker = Tracker(_opt)
    tracker.set_conn(track_queue, can_track)
    tracker.run()


if __name__ == '__main__':
    opt = create_opt(
        yolo_weights='./weights/yolov7-tiny.pt'
        , strong_sort_weights='./weights/osnet_x0_25_msmt17.pt'
        , config_strongsort='./strong_sort/configs/strong_sort.yaml'
        , project='runs/track',
        source=project_setting.camera_index
    )
    opt.view_img = True
    opt.source = '/home/cjc/Desktop/beijing.mp4'
    tracker = Tracker(opt)
    can_track = Value("i", 1)  # 是否追踪 ，0 False 1
    track_result_queue = Manager().Queue(100)
    tracker.set_conn(can_track=can_track, track_queue=track_result_queue)
    tracker.run()
